5.02. Итераторы, генераторы и контекстные менеджеры
Итераторы, генераторы и контекстные менеджеры
Итератор — это объект, представляющий поток данных, который позволяет последовательно получать элементы по одному. Он реализует протокол итерации, состоящий из двух методов:
- iter() — возвращает сам итератор (обычно self).
- next() — возвращает следующий элемент или вызывает исключение StopIteration, если элементы закончились.
Итераторы лежат в основе цикла for, выражений in, а также многих встроенных функций (sum, list, tuple и др.).
class CountDown:
def __init__(self, start):
self.current = start
def __iter__(self):
return self
def __next__(self):
if self.current < 0:
raise StopIteration
value = self.current
self.current -= 1
return value
# Использование
for n in CountDown(3):
print(n) # 3, 2, 1, 0
Здесь CountDown — это итерируемый объект, который при вызове iter() возвращает итератор (в данном случае самого себя).
Итерируемый объект — тот, у которого есть метод __iter__() (например, список, строка).
Итератор — объект, возвращаемый __iter__(), реализующий __next__(). После исчерпания итератор становится «пустым» — повторные вызовы __next__() будут бросать StopIteration.
Цикл for работает следующим образом:
- Вызывает iter(iterable) → получает итератор.
- Пока не возникнет StopIteration, вызывает next(iterator) и присваивает значение переменной.
- При завершении итерации корректно освобождает ресурсы (если нужно).
Это означает, что любой объект можно использовать в for, если он следует протоколу итерации.
Генератор — это специальный вид итератора, создаваемый с помощью функции, содержащей ключевое слово yield. Вместо возврата значения и завершения (return), yield приостанавливает выполнение функции, сохраняя текущее состояние (локальные переменные, позицию в коде), и возвращает значение. При следующем вызове __next__() выполнение продолжается с места остановки.
Генераторы реализуют ленивую загрузку (lazy evaluation) — значения генерируются по мере необходимости, а не все сразу.
def count_up_to(n):
i = 1
while i <= n:
yield i
i += 1
gen = count_up_to(3)
print(next(gen)) # 1
print(next(gen)) # 2
print(next(gen)) # 3
# next(gen) # StopIteration
Генераторы можно использовать в for:
for value in count_up_to(5):
print(value)
Преимущества ленивой загрузки - экономия памяти (не хранятся все элементы), возможность работы с бесконечными последовательностями, раннее использование данных (первый элемент доступен до завершения генерации).
Пример бесконечного генератора:
def endless_counter():
n = 0
while True:
yield n
n += 1
Генераторы списков (list comprehensions).
Синтаксис [expr for item in iterable if condition] создаёт список на основе выражения.
squares = [x**2 for x in range(5)] # [0, 1, 4, 9, 16]
evens = [x for x in range(10) if x % 2 == 0]
Все значения генераторов вычисляются сразу и хранятся в памяти, поддерживают вложенные циклы и фильтрацию, и синтаксически удобные для создания списков. При этом, не стоит использовать генераторы списков для побочных эффектов (например, вызова функций без сохранения результата) — это нарушает читаемость и может быть ошибкой.
Выражения-генераторы (generator expressions) аналогичны генераторам списков, но используют круглые скобки и возвращают генератор, а не список.
squares_gen = (x**2 for x in range(5))
print(type(squares_gen)) # <class 'generator'>
print(next(squares_gen)) # 0
Разница между [...] и (...) принципиальна:
[x**2 for x in range(1000000)]— создаст миллион объектов в памяти.(x**2 for x in range(1000000))— потребляет константный объём памяти.
Выражения-генераторы подходят для передачи в функции, ожидающие итератор:
total = sum(x**2 for x in range(10))
max_value = max(x for x in data if x > 0)
Ключевое слово yield from позволяет делегировать часть генерации другой функции-генератору.
def sub_generator():
yield "a"
yield "b"
def main_generator():
yield 1
yield from sub_generator() # вставляет значения
yield 2
for item in main_generator():
print(item) # 1, 'a', 'b', 2
yield from особенно полезен при работе с вложенными итерациями, деревьями, асинхронным кодом (где используется для await-подобного поведения).
Контекстный менеджер — это объект, определяющий поведение перед входом в блок кода и после его завершения, независимо от того, завершился ли блок нормально или с исключением.
Он реализует протокол:
__enter__()— вызывается при входе в блок with.__exit__(exc_type, exc_value, traceback)— вызывается при выходе из блока, даже если произошло исключение.
Основная цель — гарантированное освобождение ресурсов: закрытие файлов, соединений с БД, разблокировка мьютексов и т.п. Синтаксис with:
with open('file.txt') as f:
content = f.read()
# Файл автоматически закрывается здесь, даже если было исключение
Эквивалентно:
f = open('file.txt')
try:
content = f.read()
finally:
f.close()
Но with делает код более читаемым и надёжным.
Модуль itertools предоставляет набор высокоэффективных инструментов для работы с итераторами. Его функции возвращают ленивые итераторы, что делает их идеальными для обработки больших или бесконечных потоков данных.
Бесконечные итераторы:
from itertools import count, cycle, repeat
for n in count(1): # 1, 2, 3, ...
if n > 3: break
print(n)
for c in cycle('AB'): # A, B, A, B, ...
pass
Терминирующие итераторы:
from itertools import takewhile, dropwhile, chain
list(takewhile(lambda x: x < 5, [1, 4, 6, 4, 1])) # [1, 4]
list(chain([1, 2], [3, 4])) # [1, 2, 3, 4]
Комбинаторные итераторы
from itertools import product, permutations, combinations
list(product('AB', repeat=2)) # [('A','A'), ('A','B'), ...]
list(combinations('ABC', 2)) # [('A','B'), ('A','C'), ('B','C')]